from functools import cached_property
import json
import logging
import os
from concurrent import futures
from pathlib import Path
from typing import Any, Generator
from typing_extensions import Self
from pydantic import BaseModel, Field, model_validator
from pysam import FastaFile, TabixFile, VariantFile, VariantHeader, VariantRecord, tabix_compress, tabix_index  # pylint: disable=E0611
from sklearn.cluster import DBSCAN
from .sub import SUB
from .dei import DEL
from .ins import INS
from .snp import SNP

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("OpenAnno SNV")


class OpenAnno(BaseModel):
    input_file: Path = Field(alias="input")
    output_file: Path = Field(alias="output")
    genome_file: Path | None = Field(alias="genome")
    genepred_file: Path | None = Field(alias="genepred")
    filter_based_files: list[Path] = Field(alias="filter_based")
    filter_based_directories: list[Path] = Field(alias="filter_based_dir")
    region_based_files: list[Path] = Field(alias="region_based")
    gene_based_custom_file: Path | None = Field(alias="custom")
    gene_id_file: Path | None = Field(alias="gene_id")
    chromosomes: list[str] = Field(alias="chrom")
    process: int
    gene_id_dict: dict[str, str] = Field(default_factory=dict)
    gene_based_custom_dict: dict[str, dict[str, Any]] = Field(default_factory=dict)

    class Config:
        arbitrary_types_allowed = True

    @model_validator(mode="after")
    def init_data(self) -> Self:
        if self.gene_id_file:
            with open(self.gene_id_file) as reader:
                self.gene_id_dict = dict(map(lambda x: (f"{x[0]}:{x[1]}", x[2]), filter(lambda y: len(y) >= 3, map(lambda z: z.strip().split("\t"), reader))))
        if self.gene_based_custom_file:
            with open(self.gene_based_custom_file) as reader:
                self.gene_based_custom_dict = json.load(reader)
        return self

    @cached_property
    def annotation_header(self) -> VariantHeader:
        header = VariantHeader()
        if self.genepred_file:
            header.info.add(id="GENE", number=".", type="String", description="Gene Symbol")
            header.info.add(id="GENE_ID", number=".", type="String", description="Gene Entrez ID")
            header.info.add(id="REGION", number=".", type="String", description="Region in gene, eg: exonic, intronic, UTR3, UTR5")
            header.info.add(id="EVENT", number=".", type="String", description="Variant Event, eg: missense, nonsense, splicing")
            header.info.add(id="DETAIL", number=".", type="String", description="Gene detail, FORMAT=Gene:Transcript:Exon:NA_CHANGE:AA_CHANGE")
        for gene_based_file in filter(None, self.filter_based_files + list(map(lambda x: next(x.glob("*.vcf.gz"), None), self.filter_based_directories))):
            with VariantFile(str(gene_based_file)) as reader:
                for key, item in reader.header.info.items():
                    header.info.add(id=key, number=item.number, type=item.type, description=item.description)
        for region_based_file in self.region_based_files:
            with TabixFile(str(region_based_file)) as reader:
                if reader.header:
                    for line in reader.header:
                        header.add_line(line)
                else:
                    name = os.path.basename(str(self)).split(".")[0]
                    header.info.add(id=name, number=1, type="String", description=name)
        return header

    def iter_snvs(self) -> Generator[list[SNP | INS | DEL | SUB], Any, None]:
        with VariantFile(str(self.input_file)) as reader:
            for chrom in self.chromosomes:
                try:
                    records = reader.fetch(chrom)
                except:
                    logger.warning("'%s' not found when reading input: %s", chrom, self.input_file)
                else:
                    variants: list[VariantRecord]
                    intervals: list[tuple[int, int]]
                    variants, intervals = zip(*map(lambda x: (x, (x.start, x.stop)), records))
                    dbscan = DBSCAN(eps=10000, min_samples=1)
                    labels = dbscan.fit_predict(intervals)
                    snvs = {}
                    for variant, label in zip(variants, labels):
                        snv_cls = (SNP if len(variant.alts[0]) == 1 else INS) if len(variant.ref) == 1 else (DEL if len(variant.alts[0]) == 1 else SUB)
                        snvs.setdefault(label, []).append(snv_cls.from_variant_record(variant))
                    yield from snvs.values()

    def annotate_gene_based(self, snvs: list[SNP | INS | DEL | SUB]):
        genepred_reader = self.genepred_file and TabixFile(str(self.genepred_file))
        genome_reader = self.genome_file and FastaFile(str(self.genome_file))
        for snv in snvs:
            custom_info = self.gene_based_custom_dict.get(snv.pk, {})
            if not custom_info and genepred_reader and genome_reader:
                snv.annotate_gene_based(genepred_reader, genome_reader, self.gene_id_dict)
            else:
                snv.info.update(custom_info)

    def annotate_region_based(self, snvs: list[SNP | INS | DEL | SUB]):
        for region_based_file in self.region_based_files:
            with TabixFile(str(region_based_file)) as reader:
                for snv in snvs:
                    snv.annotate_region_based(reader)

    def annotate_filter_based(self, snvs: list[SNP | INS | DEL | SUB]) -> dict[str, dict[str, Any]]:
        chrom, start, end = snvs[0].chrom, snvs[0].start, snvs[-1].end
        filter_based_files = self.filter_based_files + list(map(lambda x: x / f"{chrom}.vcf.gz", self.filter_based_directories))
        results = dict(map(lambda x: (x.pk, x.info), snvs))
        for filter_based_file in filter_based_files:
            with VariantFile(str(filter_based_file)) as reader:
                for row in filter(lambda x: x.alts, reader.fetch(chrom, start, end)):
                    pk = f"{row.chrom}:{row.pos}:{row.ref}:{row.alts[0]}"
                    if results.get(pk):
                        results[pk].update(dict(row.info))
        return results

    def annotate(self, snvs: list[SNP | INS | DEL | SUB]) -> dict[str, dict[str, Any]]:
        self.annotate_gene_based(snvs)
        self.annotate_region_based(snvs)
        return self.annotate_filter_based(snvs)

    def write_output(self, results: dict[str, dict[str, Any]]):
        logger.info("writing result to %s", self.output_file)
        with VariantFile(str(self.input_file)) as reader:
            reader.header.merge(self.annotation_header)
            output_vcf = self.output_file.with_suffix("") if self.output_file.suffix == ".gz" else self.output_file
            with VariantFile(str(output_vcf), "w", header=reader.header) as writer:
                for chrom in self.chromosomes:
                    try:
                        records = reader.fetch(chrom)
                    except:
                        logger.warning(" '%s' not found when writing output: %s", chrom, self.output_file)
                    else:
                        for record in records:
                            pk = f"{record.chrom}:{record.pos}:{record.ref}:{record.alts[0]}"
                            record.info.update(results.get(pk, {}))
                            writer.write(record)
            if self.output_file.suffix == ".gz":
                tabix_compress(str(output_vcf), self.output_file, force=True)
                tabix_index(str(self.output_file), preset="vcf", force=True)

    def run(self) -> dict[str, dict[str, Any]]:
        with futures.ProcessPoolExecutor(max_workers=self.process) as executor:
            tasks = []
            for variants in self.iter_snvs():
                tasks.append(executor.submit(self.annotate, variants))
            results = {}
            for task in futures.as_completed(tasks):
                results.update(task.result())
            self.write_output(results)

    # def run(self) -> dict[str, dict[str, Any]]:
    #     for variants in self.iter_snvs():
    #         self.annotate(variants)